#include "core/memory.h"
#include "cpu/mmu.h"
#include "dev/console.h"
//物理地址分配结构
static addr_alloc_t addr_alloc;

//内核页表
static pde_t kernel_page_dir[PAGE_NUM] __attribute__((aligned(PAGE_SIZE)));

static void addr_alloc_init(addr_alloc_t* alloc,uint8_t* bits,uint32_t start,uint32_t size,uint32_t page_size)
{
    mutex_init(&alloc->mutex);
    alloc->start =start;
    alloc->size=size;
    alloc->page_size=page_size;
     
    bitmap_init(&alloc->bitmap,bits,alloc->size/page_size,0);
}


//分配物理页
static uint32_t addr_alloc_page(addr_alloc_t* alloc,int page_count)
{
    uint32_t addr=0;
    mutex_lock(&alloc->mutex);
    //页索引
    int page_index=bitmap_alloc_nbits(&alloc->bitmap,0,page_count);
    if(page_index >= 0)
    {
        //返回物理地址
        addr = alloc->start + page_index * alloc->page_size;
    }
    mutex_unlock(&alloc->mutex);
    return addr;
}

//释放物理页
static void addr_free_page(addr_alloc_t* alloc,uint32_t addr,int page_count)
{
    mutex_lock(&alloc->mutex);
    uint32_t page_index=(addr - alloc->start) / alloc->page_size;
    bitmap_set_bit(&alloc->bitmap,page_index,page_count,0);
    mutex_unlock(&alloc->mutex);
}

//可用内存尺寸
static uint32_t total_mem_size(boot_info_t* boot_info)
{
    uint32_t mem_size=0;

    for(int i=0;i<boot_info->ram_region_count;i++)
    {
        mem_size += boot_info->ram_region[i].size;
    }

    return mem_size;
}


pte_t* find_pte(pde_t* page_dir,uint32_t vaddr,int alloc)
{
    pte_t* page_table;
    //页目录表项
    pde_t* pde=page_dir + pde_index(vaddr);
    //已映射
    if(pde->present)
    {
        page_table = (pte_t*)pde_paddr(pde);
    }
    //未映射
    else
    {
        //不分配
        if(alloc == 0)
        {
            return (pte_t *)0;
        }
        //分配物理页表
        uint32_t pg_paddr = addr_alloc_page(&addr_alloc,1);
        //分配失败
        if(pg_paddr==0)
        {
            return (pte_t*)0;
        }
        //加入页目录
        pde->v= pg_paddr | PTE_P |  PTE_U | PTE_W;
        //print_log("pg_addr: %x,pde->v: %x",pg_paddr,pde->v);
        page_table=(pte_t*)pg_paddr;
        //清空页表
        kernel_memset(page_table,0,PAGE_SIZE);
    }
    //返回页表表项
    return page_table + pte_index(vaddr);
}

//显示可有内存
static void show_mem_info (boot_info_t * boot_info) 
{
    print_log("mem region:");
    for (int i = 0; i < boot_info->ram_region_count; i++) 
    {
        print_log("[%d]: 0x%x - 0x%x", i,
                    boot_info->ram_region[i].start,
                    boot_info->ram_region[i].size);
    }
}

//映射物理页与虚拟页 设置二级页表
int memory_creat_map(pde_t* page_dir,uint32_t vaddr,uint32_t paddr,int count, uint32_t perm)
{
    for(int i=0;i<count;i++)
    {
        //找到（创建）页表
        pte_t * pte=find_pte(page_dir,vaddr,1);
        //未找到
        if(pte==(pte_t*)0)
        {
            return -1;
        }
        //映射关系不存在
        ASSERT(pte->present==0);
        //建立映射关系
        pte->v = paddr | perm | PTE_P ; 
        vaddr+=PAGE_SIZE;
        paddr+=PAGE_SIZE;
    }
    return 0;
}


//创建内核页表
void  create_kernel_table()
{
    //引入OS代码区和数据区
    extern uint8_t s_text[],e_text[],s_data[];
    //定义映射关系，映射到相同地址
    static memory_map_t kernel_map[]={
        {0,s_text,0,PTE_W},
        {s_text, e_text, s_text,0},                      //代码段
        {s_data,(void*)MEM_EBDA_START,s_data,PTE_W},     //数据段
        {(void *)CONSOLE_DISP_ADDR, (void *)(CONSOLE_DISP_END),(void *)CONSOLE_DISP_ADDR, PTE_W},
        {(void*)MEM_EXT_START,(void*)MEM_EXT_END,(void*)MEM_EXT_START,PTE_W}
    };
    //清空页目录
    kernel_memset(kernel_page_dir, 0, sizeof(kernel_page_dir));
    //设置映射关系
    for(int i=0; i<sizeof(kernel_map) / sizeof(memory_map_t);i++)
    {

        memory_map_t* map= kernel_map + i;
        //向下取整对齐
        uint32_t  vstart=down2((uint32_t)map->vstart,PAGE_SIZE);
        uint32_t  vend=up2((uint32_t)map->vend,PAGE_SIZE);
        uint32_t  paddr=down2((uint32_t)(map->pstart),PAGE_SIZE);
        //内存块的页数 
        int page_count=(vend-vstart)/PAGE_SIZE;
        //映射
        memory_creat_map(kernel_page_dir,(uint32_t)vstart,(uint32_t)map->pstart,page_count,map->perm);
    }
}


//内存管理
void mermory_init(boot_info_t* boot_info)
{ 
    extern uint8_t* mem_free_start;
    //内存管理
    
    uint8_t* mem_free=(uint8_t *)&mem_free_start;
    
    show_mem_info(boot_info); 

    //1MB以上可供OS使用的内存
    uint32_t mem_up1mb_free = total_mem_size(boot_info) - MEM_EXT_START;

    //向下对齐4KB
    mem_up1mb_free=down2(mem_up1mb_free,PAGE_SIZE);
    print_log("free memory: 0x%x,size: 0x%x",MEM_EXT_START,mem_up1mb_free);
    
    //内存分配结构初始化
    addr_alloc_init(&addr_alloc,mem_free,MEM_EXT_START,mem_up1mb_free,PAGE_SIZE);
    //在OS后的内存储存位图
    mem_free+=bitmap_byte_count(addr_alloc.size/PAGE_SIZE);

    //判断是否越界
    ASSERT(mem_free<(uint8_t*)MEM_EBDA_START);

    //创建内核页表

    print_log("create_kernel_table");
    create_kernel_table();
    //写入新页表
    mmu_set_page_dir((uint32_t)kernel_page_dir); 
}   



//创建用户虚拟内存空间
uint32_t memory_create_uvm()
{
    //创建页目录
    pde_t* page_dir = (pde_t*)addr_alloc_page(&addr_alloc,1);
    //分配失败
    if(page_dir == 0)
    {
        return 0;
    }
    //分配成功

    kernel_memset((void*)page_dir,0,PAGE_SIZE);

    uint32_t user_pde_start = pde_index(MEM_TASK_BASE);
    for(int i=0;i<user_pde_start;i++)
    {
        page_dir[i].v=kernel_page_dir[i].v; //拷贝内核页表
    }
    return (uint32_t)page_dir;
}


int memory_alloc_for_page_dir(uint32_t page_dir,uint32_t vaddr,uint32_t size,int perm)
{
    uint32_t curr_vaddr = vaddr;
    int page_count = up2(size,PAGE_SIZE) / PAGE_SIZE;
    vaddr = down2(vaddr,PAGE_SIZE);

    for(int i=0;i<page_count;i++)
    {
        uint32_t paddr=addr_alloc_page(&addr_alloc,1);
        if(paddr==0)
        {
            print_log("mem alloc failed.no memory");
            return 0;
        }
        int err = memory_creat_map((pde_t *)page_dir,curr_vaddr,paddr,1,perm);
        if(err<0)
        {
            print_log("creat memory failed. err = %d",err);
            return 0;
        }
        curr_vaddr += PAGE_SIZE;
    }
    return 0;
}


int memory_alloc_page_for(uint32_t addr,uint32_t size,int perm)
{
    return memory_alloc_for_page_dir(cur_task()->tss.cr3,addr,size,perm);
}

static pde_t * current_page_dir (void) 
{
    return (pde_t *)cur_task()->tss.cr3;
}

uint32_t memory_alloc_page (void) 
{
    //内核空间虚拟地址与物理地址相同
    return addr_alloc_page(&addr_alloc, 1);
}

void memory_free_page (uint32_t addr) 
{
    if (addr < MEM_TASK_BASE) 
    {
        addr_free_page(&addr_alloc, addr, 1);
    } 
    else 
    {
        // 进程空间，还要释放页表
        pte_t * pte = find_pte(current_page_dir(), addr, 0);
        ASSERT((pte == (pte_t *)0) && pte->present);
        addr_free_page(&addr_alloc, pte_paddr(pte), 1);
        pte->v = 0;
    }
}


//释放页表
void memory_destroy_uvm(uint32_t page_dir)
{
    uint32_t user_pde_start = pde_index(MEM_TASK_BASE);
    pde_t * pde = (pde_t *)page_dir + user_pde_start;

    ASSERT(page_dir != 0);

    // 释放页表中对应的各项，不包含映射的内核页面
    for (int i = user_pde_start; i < PAGE_NUM; i++, pde++) {
        if (!pde->present) {
            continue;
        }

        // 释放页表对应的物理页 + 页表
        pte_t * pte = (pte_t *)pde_paddr(pde);
        for (int j = 0; j < PAGE_NUM; j++, pte++) {
            if (!pte->present) {
                continue;
            }

            addr_free_page(&addr_alloc, pte_paddr(pte), 1);
        }

        addr_free_page(&addr_alloc, (uint32_t)pde_paddr(pde), 1);
    }

    // 页目录表
    addr_free_page(&addr_alloc, page_dir, 1);

}

//创建页表
uint32_t memory_copy_uvm (uint32_t page_dir) 
{
    // 复制基础页表
    uint32_t to_page_dir = memory_create_uvm();
    if (to_page_dir == 0) 
    {
        goto copy_uvm_failed;
    }

    // 再复制用户空间的各项
    uint32_t user_pde_start = pde_index(MEM_TASK_BASE );
    pde_t * pde = (pde_t *)page_dir + user_pde_start;

    // 遍历用户空间页目录项
    for (int i = user_pde_start; i < PAGE_NUM; i++, pde++) 
    {
        if (!pde->present) {
            continue;
        }

        // 遍历页表
        pte_t * pte = (pte_t *)pde_paddr(pde);
        for (int j = 0; j < PAGE_NUM; j++, pte++) 
        {
            if (!pte->present) 
            {
                continue;
            }

            // 分配物理内存
            uint32_t page = addr_alloc_page(&addr_alloc, 1);
            if (page == 0) 
            {
                goto copy_uvm_failed;
            }

            // 建立映射关系
            uint32_t vaddr = (i << 22) | (j << 12);
            int err = memory_creat_map((pde_t*)to_page_dir, vaddr, page, 1, get_pte_perm(pte));
            if (err < 0) 
            {
                goto copy_uvm_failed;
            }

            // 复制内容。
            kernel_memcpy((void *)page, (void *)vaddr, PAGE_SIZE);
        }
    }
    return to_page_dir;

copy_uvm_failed:
    if (to_page_dir) {
        memory_destroy_uvm(to_page_dir);
    }
    return -1;
}


// 获取指定虚拟地址的物理地址
uint32_t memory_get_paddr (uint32_t page_dir, uint32_t vaddr) 
{
    pte_t * pte = find_pte((pde_t *)page_dir, vaddr, 0);
    if (pte == (pte_t *)0) 
    {
        return 0;
    }

    return pte_paddr(pte) + (vaddr & (PAGE_SIZE - 1));
}



//在不同的进程空间中拷贝字符串
int memory_copy_uvm_data(uint32_t to, uint32_t page_dir, uint32_t from, uint32_t size) 
{
    char *buf, *pa0;

    while(size > 0)
    {
        // 获取目标的物理地址
        uint32_t to_paddr = memory_get_paddr(page_dir, to);
        if (to_paddr == 0) 
        {
            return -1;
        }

        // 计算当前可拷贝的大小
        uint32_t offset_in_page = to_paddr & (PAGE_SIZE - 1);
        uint32_t curr_size = PAGE_SIZE - offset_in_page;
        if (curr_size > size) 
        {
            curr_size = size;       // 如果比较大，超过页边界，则只拷贝此页内
        }

        kernel_memcpy((void *)to_paddr, (void *)from, curr_size);

        size -= curr_size;
        to += curr_size;
        from += curr_size;
    }
    return 0;
}

//堆增长
char * sys_sbrk(int incr)
{
    task_t* task = cur_task();
    int pre_incr = incr;
    char* pre_heap_end=(char*)task->heap_end;
    ASSERT(incr>=0);
    if(incr==0)
    {
        return pre_heap_end;
    }

    uint32_t start=task->heap_end;
    uint32_t end=start+incr;

    int start_offset = start % PAGE_SIZE;
    if (start_offset) 
    {
        // 不超过1页，只调整
        if (start_offset + incr <= PAGE_SIZE) 
        {
            task->heap_end = end;
            return pre_heap_end;
        } else 
        {
            // 超过1页，先只调本页的
            uint32_t curr_size = PAGE_SIZE - start_offset;
            start += curr_size;
            incr -= curr_size;
        }
    }

    // 处理其余的，起始对齐的页边界的
    if (incr)
    {
        uint32_t curr_size = end - start;
        int err = memory_alloc_page_for(start, curr_size, PTE_P | PTE_U | PTE_W);
        if (err < 0) 
        {
            return (char *)-1;
        }
    }

    task->heap_end = end;
    return (char * )pre_heap_end;        
}